import { registerTrainingNodeUI } from "./training_node_ui_common.js";

registerTrainingNodeUI({
    comfyClass: "SimpleFLUXLoRATrainer",
    eventName: "flux_training_progress",
    themeColor: "#8b5cf6",
    chartTitle: "FLUX Training Loss",
    progressBarColor: "#8b5cf6",
    gradientColor: "#8b5cf660",
    legendIcon: "FLUX",
    titleFont: "bold 16px Arial"
});

// 注册FLUX训练参数节点扩展
app.registerExtension({
    name: "Comfy.FLUXTrainingParams",
    
    async beforeRegisterNodeDef(nodeType) {
        // 支持FLUX训练参数节点
        if (nodeType.comfyClass !== "FLUXTrainingParams" && 
            nodeType.comfyClass !== "SimpleFLUXTrainingParams") {
            return;
        }
        
        console.log(`FLUX训练参数节点: ${nodeType.comfyClass}`);
        
        const onDrawForeground = nodeType.prototype.onDrawForeground;
        nodeType.prototype.onDrawForeground = function(ctx) {
            if (onDrawForeground) {
                onDrawForeground.call(this, ctx);
            }
            
            // 添加FLUX标识
            ctx.textAlign = "center";
            ctx.font = "12px Arial";
            ctx.fillStyle = "#8b5cf6";
            ctx.fillText("FLUX", this.size[0] / 2, this.size[1] - 8);
        };
    }
}); 